# -*- coding: utf-8 -*-

# @Author: Chen Hua
# @Date  : 2020-07-11
# @Desc  : 中国GDP数据

import sys
import datetime
import time

import decimal
from sqlalchemy.sql.elements import and_

from db.MySqlConn import DBSession, JobConfig, engine as eg, CnGDP, IndexValue
import service.job_config_service as js
from util import DateUtil
import tushare as ts
from conf.settings import TUSHARE_TOKEN

##tushare sdk
ts.set_token(TUSHARE_TOKEN)
pro = ts.pro_api()
##链接数据库
session = DBSession()
job_name = 'stock_cn_gdp'

"""
详情：
https://tushare.pro/document/2?doc_id=227
"""

def process_gdp_api(begin_year, begin_quarter):
    for _ in range(1):  # 重试三次
        new_gdp = None
        try:
            param = str(begin_year) + 'Q' + str(begin_quarter)
            # df = pro.cn_gdp(start_q=param, end_q=param)
            df = pro.cn_gdp(start_q=param, end_q=param)
            if df.empty:
                print("没有数据返回，请检查是否出错...", str(begin_year),"-", str(begin_quarter))
            for index, row in df.iterrows():
                print(row)
                gdp = row['gdp'] * 100000000  # 亿
                ## 1季度直接*4，预计为全年，二季度*2 预计为全年， 三季度 * 4 / 3 ， 四季度已经是全年了。
                if begin_quarter == 1:
                    gdp = gdp * 4;
                elif begin_quarter == 2:
                    gdp = gdp * 2;
                elif begin_quarter == 3:
                    gdp = gdp * 4 / 3;
                new_gdp = CnGDP(dt=None, gdp=gdp)
        except Exception as e:
            print(str(e))
            time.sleep(2)
        else:
            return new_gdp


'''
根据年份和季度，获取季度的开始，结束日期
'''


def get_begin_end_dt_by_quarter(begin_year, begin_quarter):
    begin_month = None;
    end_month = None;
    end_day = None
    if begin_quarter == 1:
        begin_month = 1;
        end_month = 3;
        end_day = 31
    elif begin_quarter == 2:
        begin_month = 4;
        end_month = 6;
        end_day = 30
    elif begin_quarter == 3:
        begin_month = 7;
        end_month = 9;
        end_day = 30
    else:
        begin_month = 10;
        end_month = 12;
        end_day = 31
    return datetime.datetime(begin_year, begin_month, 1), datetime.datetime(begin_year, end_month, end_day)

'''
例如当前为5月份， 二季度的GDP还没有发布，
那么就循环往上，去一季度的数据*2，暂时当做二季度的gdp
'''
def get_mock_gdp(begin_year, begin_quarter):
    mock_ai = None
    while mock_ai is None and begin_year>2018:
        mock_ai = process_gdp_api(begin_year, begin_quarter)
        if mock_ai:
            return mock_ai;
        else:
            begin_quarter = begin_quarter - 1;  ##季度加一
            if begin_quarter < 1:
                begin_quarter = 4
                begin_year = begin_year - 1
            time.sleep(8)  # 睡眠一秒，因为接口限制 10次/分钟


def stock_cn_gdp_info_sync():
    try:
        now = datetime.datetime.now()
        end_quarter = (now.month - 1) // 3 + 1;
        end_year = now.year;

        begin_dt = datetime.datetime(2004, 1, 1)  ##默认初始从2004Q1开始
        ##从数据库获取上次已获取的数据日期
        spider_config = session.query(JobConfig).filter(JobConfig.job_name == job_name).one()
        if spider_config and spider_config.job_mode == 'inc' and spider_config.last_data_at is not None:
            begin_dt = spider_config.last_data_at
        begin_quarter = (begin_dt.month - 1) // 3 + 1;
        begin_year = begin_dt.year;

        while begin_year <= end_year:
            ai = process_gdp_api(begin_year, begin_quarter)
            print("成功爬取 ", str(begin_year), "~", str(begin_quarter), "的gdp信息...")
            quarter_begin_dt, quarter_end_dt = get_begin_end_dt_by_quarter(begin_year, begin_quarter)
            if ai is None: #查询出历史记录，临时计算为当前GDP
                ai = get_mock_gdp(begin_year,begin_quarter)
            if ai:
                # 将每个季度的一条数据，打散为每天一跳数据 ，并补充上每天的 证券总市值信息
                rs = session.query(IndexValue).filter(
                    and_(IndexValue.index_code == '1000002', IndexValue.dt >= quarter_begin_dt,
                         IndexValue.dt <= quarter_end_dt)).all()
                if rs:
                    for x in rs:
                        ai.dt = x.dt;
                        ai.mc = x.mc;  # 用流通市值计算
                        ai.cmc = x.cmc;  # 用流通市值计算
                        ai.mc_ratio = round(decimal.Decimal(ai.mc) / decimal.Decimal(ai.gdp), 3);
                        ai.cmc_ratio = round(decimal.Decimal(ai.cmc) / decimal.Decimal(ai.gdp), 3);
                        session.merge(ai)
                        session.commit()
            begin_quarter = begin_quarter + 1;  ##季度加一
            if begin_quarter > 4:
                begin_quarter = 1
                begin_year = begin_year + 1
            time.sleep(8)  # 睡眠一秒，因为接口限制 10次/分钟
            js.update_job_config(job_name, quarter_begin_dt)
        print('Stock gdp信息已保存到表 cn_gdp ...')
    except Exception as e:
        print('获取Stock gdp信息失败.')
        print(str(e))
        sys.exit(1)
    finally:
        session.close()


if __name__ == "__main__":
    stock_cn_gdp_info_sync()
